package com.ruoyi.shop.protocol.sse;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author: JinChen Zhu
 * @CreateTime: 2024-06-27  18:56
 */
@Slf4j
public class SseServer {
    /**
     * 当前连接数
     */
    private static AtomicInteger currentConnectionTotal = new AtomicInteger(0);
    /**
     * messageId的SseEmitter对象映射集
     */
    private static Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();

    /**
     * 创建sse连接
     *
     * @param messageId 连接id
     */
    public static SseEmitter createConnection(String messageId) {
        // 设置超时时间,0表示不超时
        SseEmitter sseEmitter = new SseEmitter(0L);
        try {
            sseEmitter.send(SseEmitter.event().data("重连成功！"));
        } catch (IOException e) {
            log.error("重连失败 ===> messageId:{}", messageId, e);
        }
        // 注册回调函数
        sseEmitter.onCompletion(() -> {
            log.info("连接断开 ===> messageId:{}", messageId);
        });
        sseEmitter.onTimeout(() -> {
            log.info("连接超时 ===> messageId:{}", messageId);
        });
        sseEmitter.onError((error) -> {
            log.info("连接异常 ===> messageId:{}", messageId);
        });
        sseEmitterMap.put(messageId, sseEmitter);

        // 记录连接总数,数量+1
        int count = currentConnectionTotal.incrementAndGet();
        log.info("当前连接数:{}", count);
        return sseEmitter;
    }

    /**
     * 发送消息
     *
     * @param messageId 连接id
     * @param data 消息内容
     */
    public static void sendMessage(String messageId, String data) {
        // 从sseEmitterMap中获取对应的SseEmitter对象
        SseEmitter sseEmitter = sseEmitterMap.get(messageId);
        if (sseEmitter != null) {
            try {
                // 发送消息
                //@formatter:off
                sseEmitter.send(data);
                //@formatter:on
            } catch (Exception e) {
                // 发送消息失败时记录日志
                log.error("发送消息失败", e);
            }
        } else {
            // 没有找到对应的sseEmitter时记录日志
            log.info("没有找到对应的sseEmitter");
        }
    }


    /**
     * 关闭连接
     *
     * @param messageId 连接id
     */
    public static void closeConnection(String messageId) {
        SseEmitter sseEmitter = sseEmitterMap.remove(messageId);
        if (sseEmitter != null) {
            currentConnectionTotal.decrementAndGet();
            sseEmitter.complete();
        }
    }
    /**
     * 发送广播
     */
    public static void sendBroadcast(String data) {
        for (Map.Entry<String, SseEmitter> entry : sseEmitterMap.entrySet()) {
            SseEmitter sseEmitter = entry.getValue();
            try {
                // 发送消息
                sseEmitter.send(data);
            } catch (Exception e) {
                log.error("发送消息失败", e);
            }
        }
    }

    /**
     * 获取当前连接数
     *
     * @return 当前连接数
     */
    public static int getCurrentConnectionTotal() {
        return currentConnectionTotal.get();
    }

}
